package org.atlantis.evnlib.internal;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

import org.atlantis.evnlib.internal.SetMap.Pair;

// just for single thread, not safe
public class SetMap<K, V> implements Iterable<Pair<K, V>> {

	private Map<K, Set<V>> internalMap = new HashMap<K, Set<V>>();

	public static <K1, V1> SetMap<K1, V1> newInstance() {
		return new SetMap<K1, V1>();
	}

	public void put(K key, V value) {
		Set<V> set = internalMap.get(key);
		if (set == null) {
			set = new HashSet<V>();
			internalMap.put(key, set);
		}
		set.add(value);
	}

	public Set<V> get(K key) {
		Set<V> set = internalMap.get(key);
		return set == null ? new HashSet<V>() : set;
	}

	@Override
	public Iterator<Pair<K, V>> iterator() {
		return new PairIterator<K, V>(internalMap.entrySet().iterator());
	}

	public static class Pair<K2, V2> {
		public final K2 key;
		public final Set<V2> value;

		private Pair(K2 key, Set<V2> value) {
			this.key = key;
			this.value = value;
		}
	}

	static class PairIterator<K3, V3> implements Iterator<Pair<K3, V3>> {

		private Iterator<Map.Entry<K3, Set<V3>>> iter;

		PairIterator(Iterator<Map.Entry<K3, Set<V3>>> iter) {
			this.iter = iter;
		}

		@Override
		public boolean hasNext() {
			return iter.hasNext();
		}

		@Override
		public Pair<K3, V3> next() {
			Map.Entry<K3, Set<V3>> entry = iter.next();
			return new Pair<K3, V3>(entry.getKey(), entry.getValue());
		}

		@Override
		public void remove() {
			iter.remove();
		}
	}
}
